#g-----------------------------------------------------------------
# pycparser: c-to-c.py
#
# Example of a C code generator from pycparser AST nodes, serving 
# as a simplistic translator from C to AST and back to C.
# Note: at this stage, the example is "alpha release" and considered 
# experimental. Please file any bugs you find in the Issues page on pycparser's
# website.
#
# Copyright (C) 2008-2011, Eli Bendersky
# License: BSD
#-----------------------------------------------------------------
from __future__ import print_function
import sys
import os
import wx

# This is not required if you've installed pycparser into
# your site-packages/ with setup.py
#
sys.path.insert(0, '..')

from pycparser import c_parser, c_ast, parse_file, parse_multi_files
from analyze_funcs import *
from param_tracking import *

functions={}

class Struct_Type(object):
	def __init__(self, set_name):
		self.name = set_name

class Union_Type(object):
	def __init__(self, set_name):
		self.name = set_name

class VarType(object):
	def __init__(self,set_base_type,set_indirection):
		self.base_type = set_base_type
		self.indirection = set_indirection
	def __str__(self):
		if(self.base_type == None):
			return (self.indirection * '*') + 'ANONYMOUS'
		return (self.indirection * '*') + self.base_type
	def copy(self):
		return VarType(self.base_type,self.indirection)
	def equals(self,other):
		return type(self) == type(other) and self.base_type==other.base_type and self.indirection == other.indirection


class MyDef(object):
	def __init__(self,set_type,func,set_var):
		assert(type(set_type) == VarType)
		self.type=set_type
		self.function = func
		self.parent=None
		self.children=[]
		self.var=set_var
	def get_var(self):
		if(self.var == None):
			return self
		return self.var
	def equals(self,other):
		return self.type.equals(other.type) and self.var.generate_ref_string(self) == other.var.generate_ref_string(other)
	def is_ref(self):
		return self.type.indirection > 0
	def follow_directions(self,directions):
		if(len(directions) == 0):
			return self
		next_step = directions[0]
		directions.pop(0)
		return self.children[next_step].follow_directions(directions)
		
		

class VarDef(MyDef):
	def __init__(self,set_type,set_name,func,set_var):
		if(set_var == None):
			super(VarDef,self).__init__(set_type,func,self)
		else:
			super(VarDef,self).__init__(set_type,func,set_var)
		self.name = set_name
		self.redirections=[]
	def __str__(self):
		return self.name+' : '+str(self.type)
	def show(self,unused_name=''):
		result = self.name +' : '+str(self.type)
		if(self.parent != None):
			result = result+'\n'+'&('+self.name+')'+' : '+str(self.parent.type)
		for child in self.children:
			result = result +'\n'+child.show('('+self.name+')')
		return result
	def copy_for_struct(self,var):
#		var = VarDef(self.type.copy(), None,self.name,self.var)
		tmp = StructDef(self.type.copy(), None,self.name,var)
		return tmp
	def copy(self):
#		var = VarDef(self.type.copy(), None,self.name,self.var)
		var = StructDef(self.type.copy(), None,self.name,self.var)
		return var
	def equals(self,other):
		return type(other) == type(self) and self.name==other.name and super(VarDef,self).equals(other)
	def generate_ref_string(self,target,base_string=''):
		if(target == (self)):
			return '('+self.name+')'
		if(target == self.parent):
			return '&('+self.name+')'
		for child in self.children:
			s = child.generate_ref_string(target,'('+self.name+')')
			if(s != None):
				return s
		return None

	def get_directions_to_var(self,target):
		if(target == self):
			return [ ]
		if(target == self.parent):
			return [ -1 ]
		i = 0
		for child in self.children:
			result = child.get_directions_to_var(target)
			if(result != None):
				return [ i ] + result
			i = i + 1
		return None

class RefDef(MyDef):
	def __init__(self,set_type,func,set_var):
		super(RefDef,self).__init__(set_type,func,set_var)
	def show(self,name=''):
		result=name+'[] : '+str(self.type)
		for child in self.children:
			result = result + '\n' + child.show('('+name+'[])')
		return result
	def copy_for_struct(self,var):
		tmp = RefDef(self.type.copy(), None,var)
		return tmp
	def copy(self):
		var = RefDef(self.type.copy(), None,self.var)
		return var
	def equals(self,other):
		return type(other) == type(self) and super(RefDef,self).equals(other)
	def generate_ref_string(self,target,base_string):
		if(target == (self)):
			return '('+base_string+'[])'
		if(target == self.parent):
			return base_string
		for child in self.children:
			s = child.generate_ref_string(target,'('+base_string+'[])')
			if(s != None):
				return s
		return None
	def get_directions_to_var(self,target):
		if(target == self):
			return [ ]
		if(target == self.parent):
			return [ -1 ]
		i = 0
		for child in self.children:
			result = child.get_directions_to_var(target)
			if(result != None):
				return [ i ] + result
			i = i + 1
		return None

class StructDef(MyDef):
	def __init__(self,set_type,func,set_field,set_var):
		super(StructDef,self).__init__(set_type,func,set_var)
		self.field_name = set_field
	def show(self,name=''):
		result = name+'.'+self.field_name + ' : '+str(self.type)
		if(self.parent != None):
			result = result + '\n'+ '&('+name+'.'+self.field_name+') : '+str(self.parent.type)
		for child in self.children:
			result = result + '\n'+child.show('('+name+'.'+self.field_name+')')
		return result
	def copy_for_struct(self,var):
		tmp = StructDef(self.type.copy(), None,self.field_name,var)
		return tmp
	def copy(self):
		var = StructDef(self.type.copy(), None,self.field_name,self.var)
		return var
	def equals(self,other):
		return type(self) == type(other) and self.field_name == other.field_name and super(StructDef,self).equals(other)
	def generate_ref_string(self, target,base_string):
		if(target == (self)):
			return '('+base_string+'.'+self.field_name+')'
		if(target == self.parent):
			return '&('+base_string+'.'+self.field_name+')'
		for child in self.children:
			s = child.generate_ref_string(target,'('+base_string+'.'+self.field_name+')')
			if(s != None):
				return s
		return None
	def get_directions_to_var(self,target):
		if(target == self):
			return [ ]
		if(target == self.parent):
			return [ -1 ]
		i = 0
		for child in self.children:
			result = child.get_directions_to_var(target)
			if(result != None):
				return [ i ] + result
			i = i + 1
		return None
		



def insert_indent(indent,s):
	lines=s.split('\n')
	result = ''
	for l in lines:
		result = result + (' '*indent)+l+'\n'
	return result


class CGenerator(object):
	""" Uses the same visitor pattern as c_ast.NodeVisitor, but modified to
	return a value from each visit method, using string accumulation in 
	generic_visit.
	"""
	def __init__(self):
        
        # Statements start with indentation of self.indent_level spaces, using
        # the _make_indent method
        #
		self.indent_level = 0
		self.current_type_def = []
		self.current_function=None
		self.current_type=[]
		self.current_typedef=None
		self.current_funcdef=None
		self.globals={}
		self.struct_defs={}
		self.union_defs={}
		self.current_vardecl=None
		self.running_dependency_step=False
		self.running_struct_step = False
		self.list_of_vars = []

	def get_current_function(self):
		if(self.current_function == None):
			return None
		else:
			return functions[self.current_function]

	def add_to_appropriate_scope(self,var):
		assert(type(var) == VarDef)

		self.list_of_vars.append(var)

		if(self.current_funcdef):
			# declaration of parameter
			functions[self.current_function].params[var.name] = var
			functions[self.current_function].params_list.append(var.name)
		
		elif(self.current_function):
			# local declaration
			functions[self.current_function].locals[var.name] = var
		elif(len(self.current_type) > 0):
			curr_t = self.current_type[len(self.current_type)-1]
			if(type(curr_t) == Struct_Type):
				if(not (curr_t.name in self.struct_defs.keys()) and not (curr_t.name in self.union_defs.keys())):
					self.struct_defs[curr_t.name]=[]
				if(curr_t.name in self.struct_defs.keys()):
					self.struct_defs[curr_t.name].append(var)
			elif(type(curr_t) == Union_Type):
				if(not (curr_t.name in self.union_defs.keys()) and not (curr_t.name in self.struct_defs.keys()) ):
					self.union_defs[curr_t.name]=[]
				if(curr_t.name in self.union_defs.keys()):
					self.union_defs[curr_t.name].append(var)
		else:
			# global declaration of variable
			self.globals[var.name]=var

	def copy_struct_helper(self,var,parent,new_parent,base_var):
		new_var = var.copy_for_struct(base_var)

		if(parent != None and parent.equals(var.parent)):
			new_var.parent = new_parent
		else:
			new_var.parent = var.parent.copy()
		if(type(new_var) == StructDef):
			new_var.parent.children.append(new_var)
			new_var.parent.var = base_var
		for child in var.children:
			new_var.children.append(self.copy_struct_helper(child,var,new_var,base_var))
		return new_var
			

	def copy_struct(self,struct_name,base_var):
		assert(struct_name in self.struct_defs)
		struct_def = self.struct_defs[struct_name]
		new_struct_def=[]
		for field in struct_def:
			assert(type(field) == VarDef)
			new_struct_def.append(self.copy_struct_helper(field,None,None,base_var))
#		for field in new_struct_def:
#			field.parent.children.append(field)
		return new_struct_def
		
    
	def visit(self, node):
		method = 'visit_' + node.__class__.__name__
		getattr(self, method, self.generic_visit)(node)
    
	def generic_visit(self, node):
        #~ print('generic:', type(node))
		if not (node is None): 
			for c in node.children():
				self.visit(c)
    
	def visit_Constant(self, n):
		pass
        
	def visit_ID(self, n):
		pass

	def visit_ArrayRef(self, n):
		self.visit(n.name)
		self.visit(n.subscript)

	def visit_StructRef(self, n):

		self.visit(n.name)
		self.visit(n.field)

	def visit_FuncCall(self, n):
		self.visit(n.name)
		self.visit(n.args)

	def visit_UnaryOp(self, n):
		self.visit(n.expr)

	def visit_BinaryOp(self, n):
		self.visit(n.left)
		self.visit(n.right)
    
	def visit_Assignment(self, n):

		if(self.current_function==None):
			print("Error processing assignment. We don't seem to be in a function right now!")
			exit()
		self.visit(n.rvalue) 
		self.visit(n.lvalue)
    
	def visit_IdentifierType(self, n):
		pass

	def generate_all_possible_references(self,varname,n,ptrs_passed):

		if(n == None): 
			return
		typ = type(n)

		if(typ == c_ast.TypeDecl):
			if(type(n.type) == c_ast.IdentifierType):
				base_type = n.type.names[0]
			elif(type(n.type) == c_ast.Struct):
				base_type = n.type.name
			elif(type(n.type) == c_ast.Union):
				base_type = varname
			else:
				print('unhandled type in typedecl in generate_all_possible_references: '+varname+' '+str(type(n.type)))
				exit()
			

			var = VarDef(VarType(base_type,ptrs_passed), n.declname, self.get_current_function(),None)
			var.var = var
			address_of = RefDef(VarType(base_type,ptrs_passed+1),self.get_current_function(),var)

			address_of.children.append(var)
			var.parent = address_of
			self.add_to_appropriate_scope(var)
			return var

		elif(typ==c_ast.PtrDecl):
			parent = self.generate_all_possible_references(varname,n.type,ptrs_passed+1)
			var = RefDef(VarType(parent.type.base_type,ptrs_passed), self.get_current_function(),parent.var)
			var.parent = parent
			var.children = parent.children
			parent.children = []
			parent.children.append(var)
			for child in var.children:
				child.parent = var
			return var
		elif(typ==c_ast.ArrayDecl):
			parent = self.generate_all_possible_references(varname,n.type,ptrs_passed+1)
			var = RefDef(VarType(parent.type.base_type,ptrs_passed), self.get_current_function(),parent.var)
			var.parent = parent
			var.children = parent.children
			parent.children = []
			parent.children.append(var)
			for child in var.children:
				child.parent = var
			return var
		else:
			print('unhandled type '+str(typ))
			exit()

	def visit_Decl(self, n, no_type=False):
		typ = type(n.type)
		variable_name = n.name

		#TODO variable name obfuscation
		if(self.running_dependency_step and (typ == c_ast.TypeDecl or typ == c_ast.PtrDecl or typ == c_ast.ArrayDecl) and not self.running_struct_step):
			self.generate_all_possible_references(variable_name,n.type,0)
		elif(self.running_struct_step and len(self.current_type) > 0):
			self.generate_all_possible_references(variable_name,n.type,0)
#		elif(self.running_struct_step and self.current_union != None):
#			self.generate_all_possible_references(variable_name,n.type,0)
		

		self._generate_decl(n)
		
		if n.bitsize: self.visit(n.bitsize)
		if n.init:
			self.visit(n.init)
    
	def visit_DeclList(self, n):
		self.visit(n.decls[0])
		if len(n.decls) > 1:
			for decl in n.decls[1:]:
				self.visit_Decl(decl,no_type=True)
    
	def visit_Typedef(self, n):
		self.current_typedef=n.name
		self._generate_type(n.type)
		self.current_typedef=None
    
	def visit_Cast(self, n):
		self._generate_type(n.to_type)
		self.visit(n.expr)
    
	def visit_ExprList(self, n):
		for expr in n.exprs:
			self.visit(expr)
    
	def visit_Enum(self, n):
		if n.values:
			for i, enumerator in enumerate(n.values.enumerators):
				if enumerator.value: 
					self.visit(enumerator.value)
    
#	def visit_Struct(self, n):
#		if n.decls:
#			for decl in n.decls:
#				self.visit(decl) 
        
	def visit_FuncDef(self, n):
		if(not n.decl.name in functions):
			functions[n.decl.name]=Function(n.decl.name)
		self.current_funcdef = n.decl.name
	
		self.current_function=n.decl.name
		self.visit(n.decl)
		self.current_funcdef=None
		self.visit(n.body)
		self.current_function=None

	def visit_FileAST(self, n):
		for ext in n.ext:
			self.visit(ext)

	def visit_Compound(self, n):
		if(n.block_items != None):
			for stmt in n.block_items:
				self._generate_stmt(stmt)
    
	def visit_ParamList(self, n):
		for param in n.params:
			self.visit(param)

	def visit_Return(self, n):
		if n.expr: self.visit(n.expr)

	def visit_Break(self, n):
		pass
        
	def visit_Continue(self, n):
		pass
    
	def visit_TernaryOp(self, n):
		self.visit(n.cond)
		self.visit(n.iftrue)
		self.visit(n.iffalse)
    
	def visit_If(self, n):
		if n.cond: self.visit(n.cond)
		self._generate_stmt(n.iftrue, add_indent=True)
		if n.iffalse: 
			self._generate_stmt(n.iffalse, add_indent=True)
    
	def visit_For(self, n):
		if n.init: self.visit(n.init)
		if n.cond: self.visit(n.cond)
		if n.next: self.visit(n.next)
		self._generate_stmt(n.stmt, add_indent=True)

	def visit_While(self, n):
		if n.cond: self.visit(n.cond)
		self._generate_stmt(n.stmt, add_indent=True)

	def visit_DoWhile(self, n):
		self._generate_stmt(n.stmt, add_indent=True)
		if n.cond: self.visit(n.cond)

	def visit_Switch(self, n):
		self.visit(n.cond)
		self._generate_stmt(n.stmt, add_indent=True)
    
	def visit_Case(self, n):
		self.visit(n.expr)
		self._generate_stmt(n.stmt, add_indent=True)
    
	def visit_Default(self, n):
		self._generate_stmt(n.stmt, add_indent=True)

	def visit_Label(self, n):
		self._generate_stmt(n.stmt)

	def visit_Goto(self, n):
		pass

	def visit_EllipsisParam(self, n):
		pass

	def visit_Struct(self, n):
		if(self.running_struct_step):
			if(n.name==None):
				#self.current_struct=self.current_typedef
				self.current_type.append(Struct_Type(self.current_typedef))
			else:
				#self.current_struct=n.name
				self.current_type.append(Struct_Type(n.name))
		
			self._generate_struct_union(n, 'struct')
			#self.current_struct = None
			self.current_type.pop()

	def visit_Typename(self, n):
		self._generate_type(n.type)
        
	def visit_Union(self, n):
		if(self.running_struct_step):
			if(n.name == None):
				self.current_type.append(Union_Type(self.current_typedef))
				#self.current_union = self.current_typedef
			else:
				#self.current_union = n.name
				self.current_type.append(Union_Type(n.name))
			self._generate_struct_union(n, 'union')
			#self.current_union = None
			self.current_type.pop()
		

	def visit_NamedInitializer(self, n):
		self.visit(n.expr)

	def _generate_struct_union(self, n, name):
		""" Generates code for structs and unions. name should be either 
		'struct' or union.
		"""
		if n.decls:
			for decl in n.decls:
				self._generate_stmt(decl)

	def _generate_stmt(self, n, add_indent=False):
		""" Generation from a statement node. This method exists as a wrapper
		for individual visit_* methods to handle different treatment of 
		some statements in this context.
		"""
		typ = type(n)
		self.visit(n)
        
	def _generate_decl(self, n):
		if n.funcspec: n.funcspec
		if n.storage: n.storage
		self._generate_type(n.type)
   
	def _generate_type(self, n, modifiers=[]):
		typ = type(n)
		
        
		if typ == c_ast.TypeDecl:
			self.visit(n.type)
            
			for i, modifier in enumerate(modifiers):
				if isinstance(modifier, c_ast.ArrayDecl):
					self.visit(modifier.dim)
				elif isinstance(modifier, c_ast.FuncDecl):
					self.visit(modifier.args)
		elif typ == c_ast.Decl:
			self._generate_decl(n.type)
		elif typ == c_ast.Typename:
			self._generate_type(n.type)
		elif typ in (c_ast.ArrayDecl, c_ast.PtrDecl, c_ast.FuncDecl):
			self._generate_type(n.type, modifiers + [n])
		else:
			self.visit(n)

def get_equiv(current, redirection, gen, pairings):
	pairings.append( [ redirection, current] )
	v = redirection.var
	for r in v.redirections:
		direct = v.get_directions_to_var(r[0])
		tmp = current.follow_directions(direct)
		get_equiv(tmp, r[1], gen, pairings)

def eval_redirections(gen):
	pairings = [ ]
	for g in gen.globals:
		definition = gen.globals[g]
		for r in definition.redirections:
			direct = definition.get_directions_to_var(r[0])
			tmp = definition.follow_directions(direct)
			equivalencies = get_equiv(tmp, r[1], gen, pairings)
	for f in gen.functions:
		func = gen.functions[f]
		for l in func.locals:
			definition = func.locals[l]
			for r in definition.redirections:
				direct = definition.get_directions_to_var(r[0])
				tmp = definition.follow_directions(direct)
				equivalencies = get_equiv(tmp,r[1],gen,pairings)
	return pairings

			
			

def def_in_list(d,l):
	for element in l:
		if(element.equals(d)):
			return True
	return False

def do_analysis(files, include_paths,defines):
	generator = CGenerator()
	ast = parse_multi_files(files,include_paths,defines, os.path.abspath(__file__)[0:len(os.path.abspath(__file__))-10],True)

	generator.running_struct_step = True
	generator.visit(ast)
	generator.running_struct_step = False
	generator.running_dependency_step=True
	generator.visit(ast)

	generator.running_dependency_step = False

		# link references to structures to their definitions
	i = 0
	while(i < len(generator.list_of_vars)):
		var = generator.list_of_vars[i]
		if not var.type.base_type in generator.struct_defs:
			generator.list_of_vars.remove(var)
		else:
			i = i + 1

	for var in generator.list_of_vars:
		new_struct_def = generator.copy_struct(var.type.base_type,var)
		tmp = var
		while(len(tmp.children) > 0):
			assert(len(var.children) == 1)
			tmp = tmp.children[0]
#		print('appending to '+tmp.show())
		for field in new_struct_def:
			tmp.children.append(field)
	
	new_generator=FunctionAnalyzer(functions,generator.globals,generator.struct_defs)
	new_generator.visit(ast)
	param_tracker = ParamTracker(functions,new_generator.globals,new_generator.struct_defs)
	param_tracker.visit(ast)
	equivs = eval_redirections(param_tracker)

	redirection_buckets = { }
	retain_ordering = [ ]
	if(len(equivs) > 0):
		print('equivalencies:')
		for n in equivs:
			if(not n[1] in redirection_buckets):
				redirection_buckets[n[1]] = [ ]
				retain_ordering.append(n[1])
			redirection_buckets[n[1]].append(n[0])

		for var in retain_ordering:
			if(var.function == None):
				key='global : '+var.var.generate_ref_string(var)
			else:
				key=var.function.name+' : '+var.var.generate_ref_string(var)
			print('    '+key)
			for ref in redirection_buckets[var]:
				print('        '+ref.function.name+' : '+ref.var.generate_ref_string(ref))
	
		print()

	if(len(generator.globals) > 0):
		print('globals:')
		for g in generator.globals:
			print(str(generator.globals[g]))
		print()
	if(len(generator.struct_defs) > 0):
		print('struct definitions:')
		for s in generator.struct_defs:
			print('    '+s+':')
			for field in generator.struct_defs[s]:
				print('        '+str(field))
		print()
	if(len(generator.union_defs) > 0):
		print('union definitions:')
		for s in generator.union_defs:
			print('    '+s+':')
			for field in generator.union_defs[s]:
				print('        '+str(field))
		print()




def parse_paths(concat_paths):
	if(concat_paths.find(':') == -1):
		return [ concat_paths ]
	return concat_paths.split(':')
	

#------------------------------------------------------------------------------
if __name__ == "__main__":
	if(len(sys.argv) >= 2):
		files=[]
		include_path=[]
		defines=[]
		i = 1
		while(i < len(sys.argv)):
			arg = sys.argv[i]
			if(arg == '-I'):
				if(len(sys.argv) < i+2):
					print('Missing path after -I command line argument')
					quit()
				paths = parse_paths(sys.argv[i+1])
				for p in paths:
					include_path.append(p)
				i = i + 1
			elif(arg[0:2] == '-I'):
				paths = parse_paths(arg[2:len(arg)])
				for p in paths:
					include_path.append(p)
			elif(arg == '-D'):
				if(len(sys.argv) < i+2):
					print('Missing path after -D command line argument')
					quit()
				defines.append(sys.argv[i+1])
				i = i + 1
			elif(arg[0:2] == '-D'):
				defines.append(arg[2:len(arg)])
			else:
				files.append(arg)
			i = i + 1

		do_analysis(files, include_path,defines)
		
		for func in functions.keys():
			if(functions[func].isempty()):
				continue
			print('function '+str(func)+':')
			if(len(functions[func].reads) > 0):
				print('    reads:'+str(len(functions[func].reads)))
				for r in functions[func].reads:
					print('        '+str(r[1].var.var.generate_ref_string(r[1].var))+', '+r[0])
			if(len(functions[func].writes) > 0):
				print('    writes:'+str(len(functions[func].writes)))
				for w in functions[func].writes:
					print('        '+str(w[1].var.var.generate_ref_string(w[1].var))+', '+w[0])
			if(len(functions[func].possible_writes) > 0):
				print('    possibles:'+str(len(functions[func].possible_writes)))
				for w in functions[func].possible_writes:
					print('        '+str(w[1].var.var.generate_ref_string(w[1].var))+', '+w[0])
			print()
			displayed = []
			if(len(functions[func].calls) > 0):
				print('    calls:'+str(len(functions[func].calls)))
				for f in functions[func].calls:
					if(not f.name in displayed):
						displayed.append(f.name)
						print('        '+f.name)
			displayed=[]
			if(len(functions[func].called_by) > 0):
				print('    called by:'+str(len(functions[func].called_by)))
				for f in functions[func].called_by:
					if(not f.func.name in displayed):
						displayed.append(f.func.name)
						print('        '+f.func.name)
		
			if(len(functions[func].calls) > 0 or len(functions[func].called_by) > 0):
				print()

	else:
		print("Please provide a filename as argument")

